#!/usr/bin/env python
#
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging
import subprocess
import argparse
import os
import sys

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

sys.path.append(
    os.path.join(SCRIPT_DIR, os.pardir, os.pardir, os.pardir, os.pardir,
                 'third_party', 'perfetto', 'python'))

sys.path.append(
    os.path.join(SCRIPT_DIR, os.pardir, os.pardir, os.pardir, 'mac', 'power',
                 'protos', 'third_party', 'pprof'))

PERFETTO_DIR = os.path.normpath(
    os.path.join(os.path.dirname(__file__), os.pardir, os.pardir, os.pardir,
                 os.pardir, 'third_party', 'perfetto'))

TRACECONV_PATH = os.path.join(PERFETTO_DIR, 'tools', 'traceconv')

ARTIFACTS_DIR = os.path.normpath(
    os.path.join(os.path.dirname(__file__), os.pardir, os.pardir, 'artifacts'))

_DESCRIPTION = ("""
Symbolizes and deobfuscates the traces generated by a run of the benchmark for
the contrib power stories. This script will automatically find all the runs and
only generate results if needed.
""")

_USAGE = """
process_results --out-dir=<out/Release>
"""

_POWER_STORIES = ['power_scroll_top']


def _CreateArgumentParser():
  parser = argparse.ArgumentParser(description=_DESCRIPTION, usage=_USAGE)

  parser.add_argument('--out-dir', action='store', dest='out_dir')
  parser.add_argument('-v',
                      '--verbose',
                      action='count',
                      dest='verbosity',
                      default=0,
                      help='Increase verbosity level (repeat as needed)')

  return parser


def _GetTraceEventProto(story_results_path):
  for f in os.scandir(os.path.join(story_results_path, 'trace', 'traceEvents')):
    if not f.is_file():
      continue
    if f.name.endswith('.pb') and f.name.count('.') == 1:
      return f.path
  return None


def _GetProcessNameMapping(trace_path):
  # We can not have it at the top level as perfetto is not checked out in some
  # bots and this file is loaded in some tests to look for benchmarks
  from perfetto.trace_processor import TraceProcessor  # pylint: disable=import-error,import-outside-toplevel
  mappings = {}
  logging.info("Getting prcess mappings from: %s ",
               trace_path[len(ARTIFACTS_DIR) + 1:])
  tp = TraceProcessor(file_path=trace_path)
  res = tp.query('SELECT pid, name FROM process')
  for r in res:
    mappings[r.pid] = (r.name if r.name is not None else 'null')
  return mappings


def _AddProcessNameToProfile(input_path, output_path, process_name):
  # We can not have it at the top level as perfetto is not checked out in some
  # bots and this file is loaded in some tests to look for benchmarks
  import profile_pb2  # pylint: disable=import-error,import-outside-toplevel
  profile = profile_pb2.Profile()
  with open(input_path, "rb") as f:
    profile.ParseFromString(f.read())

  process_name_id = len(profile.string_table)
  profile.string_table.append(process_name)

  new_function = profile_pb2.Function()
  new_function.id = max([a.id for a in profile.function]) + 1
  new_function.name = process_name_id
  profile.function.append(new_function)

  new_location = profile_pb2.Location()
  new_location.id = max([a.id for a in profile.location]) + 1
  new_line = new_location.line.add()
  new_line.function_id = new_function.id
  profile.location.append(new_location)

  for sample in profile.sample:
    sample.location_id.append(new_location.id)

  with open(output_path, "wb") as f:
    f.write(profile.SerializeToString())


def _ProcessProfile(input_path, output_path, trace_path):
  logging.info('Processing profile: %s', input_path[len(ARTIFACTS_DIR) + 1:])
  pid_map = _GetProcessNameMapping(trace_path)

  if not os.path.exists(output_path):
    os.mkdir(output_path)

  for p in os.scandir(input_path):
    processed_profile_path = os.path.join(output_path, p.name)
    parts = p.name.split('.')
    if (len(parts) != 5 or not parts[3].isnumeric() or len(parts[3]) == 0):
      logging.warning('Unexpected profile file: %s', p.name)
      continue
    pid = int(parts[3])
    _AddProcessNameToProfile(p.path, processed_profile_path,
                             pid_map.get(pid, '(unknown)'))


def _RunTraceConv(command, build_out_path):
  traceconv_env = os.environ.copy()
  traceconv_env['PERFETTO_SYMBOLIZER_MODE'] = 'index'
  traceconv_env['PERFETTO_BINARY_PATH'] = os.path.join(build_out_path,
                                                       'lib.unstripped')
  traceconv_env['PERFETTO_PROGUARD_MAP'] = 'org.chromium.*=' + os.path.join(
      build_out_path, 'apks', 'ChromePublic.apk.mapping')

  return subprocess.run([TRACECONV_PATH] + command,
                        env=traceconv_env,
                        check=True,
                        capture_output=True)


def _Symbolize(input_path, output_path, build_out_path):
  logging.info('Symbolizing: %s', input_path[len(ARTIFACTS_DIR) + 1:])
  cmd = ['symbolize', input_path, output_path]
  return _RunTraceConv(cmd, build_out_path)


def _Deobfuscate(input_path, output_path, build_out_path):
  logging.info('Deobfuscating: %s', input_path[len(ARTIFACTS_DIR) + 1:])
  cmd = ['deobfuscate', input_path, output_path]
  return _RunTraceConv(cmd, build_out_path)


def _Profile(input_path, output_path, build_out_path):
  logging.info('Profiling: %s', input_path[len(ARTIFACTS_DIR) + 1:])
  cmd = ['profile', "--perf", input_path]
  run = _RunTraceConv(cmd, build_out_path)
  if len(run.stdout) == 0:
    os.mkdir(output_path)
    return
  profile_out_dir = run.stdout.split()[-1]
  os.rename(profile_out_dir, output_path)


def _Cat(input_paths, out_path):
  logging.info('Concatenating: %s', ", ".join(input_paths))
  with open(out_path, "wb") as out_f:
    for input_path in input_paths:
      with open(input_path, "rb") as in_f:
        out_f.write(in_f.read())


def _ProcessStory(story_results_path, build_out_path):
  if not _HasTraceEvents(story_results_path):
    logging.info('Skipping results with no traceEvents: %s',
                 story_results_path[len(ARTIFACTS_DIR) + 1:])
    return
  logging.info('Processing results: %s',
               story_results_path[len(ARTIFACTS_DIR) + 1:])

  traceconv_env = os.environ.copy()
  traceconv_env['PERFETTO_SYMBOLIZER_MODE'] = 'index'
  traceconv_env['PERFETTO_BINARY_PATH'] = os.path.join(build_out_path,
                                                       'lib.unstripped')
  traceconv_env['PERFETTO_PROGUARD_MAP'] = 'org.chromium.*=' + os.path.join(
      build_out_path, 'apks', 'ChromePublic.apk.mapping')

  trace_file = _GetTraceEventProto(story_results_path)
  if trace_file is None:
    raise Exception('Did not find trace file in %s' % story_results_path)

  base_path = os.path.splitext(trace_file)[0]
  symbols_path = base_path + '.symbols.pb'
  mappings_path = base_path + '.mappings.pb'
  combined_path = base_path + '.combined.pb'
  profile_path = os.path.join(os.path.dirname(combined_path), 'profile')
  processed_profile_path = os.path.join(os.path.dirname(combined_path),
                                        'processed_profile')

  if not os.path.isfile(symbols_path):
    _Symbolize(trace_file, symbols_path, build_out_path)

  if not os.path.isfile(mappings_path):
    _Deobfuscate(trace_file, mappings_path, build_out_path)

  if not os.path.isfile(combined_path):
    _Cat([trace_file, symbols_path, mappings_path], combined_path)

  if not os.path.isdir(profile_path):
    _Profile(combined_path, profile_path, build_out_path)

  if not os.path.isdir(processed_profile_path):
    _ProcessProfile(profile_path, processed_profile_path, trace_file)


def _IsPowerStory(name):
  return name.startswith("contrib_power_mobile")


def _IterateRunPaths():
  for r in os.scandir(ARTIFACTS_DIR):
    if not r.is_dir() or not r.name.startswith('run'):
      continue
    yield r.path


def _IterateStoryPaths(run_path):
  for s in os.scandir(run_path):
    if s.is_dir() and _IsPowerStory(s.name):
      yield s.path


def _HasTraceEvents(story_results_path):
  return os.path.isdir(os.path.join(story_results_path, 'trace', 'traceEvents'))


def _ProcessRun(run_path):
  combined_profile_path = os.path.join(run_path, 'combined_profile.pb.gz')
  if os.path.exists(combined_profile_path):
    return
  logging.info('Generating combined profile: %s',
               combined_profile_path[len(ARTIFACTS_DIR) + 1:])
  profiles = []
  for s in _IterateStoryPaths(run_path):
    processed_profile_path = os.path.join(s, 'trace', 'traceEvents',
                                          'processed_profile')
    if not os.path.isdir(processed_profile_path):
      continue
    for p in os.scandir(processed_profile_path):
      profiles.append(p.path)
  if len(profiles) == 0:
    return
  subprocess.run(['pprof', '-proto', '-output', combined_profile_path] +
                 profiles,
                 check=True,
                 capture_output=True)


def main():
  parser = _CreateArgumentParser()
  args = parser.parse_args()

  if args.verbosity >= 2:
    logging.getLogger().setLevel(logging.DEBUG)
  elif args.verbosity == 1:
    logging.getLogger().setLevel(logging.INFO)
  else:
    logging.getLogger().setLevel(logging.WARNING)

  if not args.out_dir:
    raise Exception('--out_dir missing')

  build_out_path = os.path.abspath(os.path.expanduser(args.out_dir))

  if not os.path.isdir(build_out_path):
    raise Exception('Unable to find out_dir %s' % build_out_path)

  for run_path in _IterateRunPaths():
    for s in _IterateStoryPaths(run_path):
      _ProcessStory(s, build_out_path)
    _ProcessRun(run_path)


if __name__ == '__main__':
  sys.exit(main())
